'''
Author: daniel
Date: 2024-03-23 12:54:10
LastEditTime: 2024-03-23 12:56:32
LastEditors: daniel
Description: 
FilePath: /matrixVT/scripts/train.py
have a nice day
'''
import yaml
import subprocess
import os
from IPython import embed


def main():
    # Load config
    config_path = 'scripts/train.yaml'
    with open(config_path, 'r') as file:
        config = yaml.safe_load(file)

    # Data convert
    data_root = config['data_prepare']['data_root']
    save_name = config['data_prepare']['save_name']
    clip_num = config['data_prepare']['clip_num']
    train_num = config['data_prepare']['train_num']


    work_dir = 'data_convert'
    command = [
        'python', 'convert.py',
        '--data_path', data_root,
        '--clip_num', str(clip_num),
        '--save_path', f'../data/{save_name}'
    ]
    ret = subprocess.run(command, cwd=work_dir)
    if ret.returncode == 0:
        print("Success: Converted Changan to NuScenes Format")
    else:
        print("Error:", ret)
    work_dir = '.'
    command = [
        'python', 'tools/gen_info.py',
        '--train_num', str(train_num),
        '--save_name', save_name
    ]
    ret = subprocess.run(command, cwd=work_dir)
    if ret.returncode == 0:
        print("Success: Converted NuScenes to pkl Format")
    else:
        print("Error:", ret)
        
    # # Training
    data_name = config['train']['data_name']
    exp_name = config['train']['exp_name']
    batch_size = config['train']['batch_size']
    gpu_num = config['train']['gpu_num']
    model = config['train']['model']
    
    if model == 'MatrixVT':
        work_dir = '.'
        command = [
            'python', 'bevdepth/exps/nuscenes/MatrixVT/matrixvt_bev_depth_lss_r50_256x704_128x128_24e_ema.py',
            '--exp', exp_name,
            '--data_root', data_name,
            '--amp_backend', 'native',
            '-b', str(batch_size),
            '--gpus', str(gpu_num)
        ]
        ret = subprocess.run(command, cwd=work_dir)
        if ret.returncode == 0:
            print("Finish Training!")
        else:
            print("Error:", ret)
    else:
        raise ValueError(f'Model {model} not supported yet!')

if __name__ == '__main__':
    main()
